package cn.bellychang.wangzheng.greedy;

import java.util.*;

public class HuffmanSolution {

    public EncodeWrapper encode(String source) {
        char[] chars = source.toCharArray();
        HashMap<Character, Integer> map = new HashMap<>(chars.length);
        for (int i = 0; i < chars.length; i++) {
            if (!map.containsKey(chars[i])) {
                map.put(chars[i], 1);
            } else {
                map.put(chars[i], map.get(chars[i]) + 1);
            }
        }
        Map<Character, Integer> linkedMap = MapUtil.sortByValue(map);
        HashMap<Character, String> encodeMap = new HashMap<>(chars.length);
        HashMap<String, Character> decodeMap = new HashMap<>(chars.length);
        int i = 0;
        for (Map.Entry<Character, Integer> entry : linkedMap.entrySet()) {
            encodeMap.put(entry.getKey(), generateString(i, linkedMap.size()));
            decodeMap.put(generateString(i, linkedMap.size()), entry.getKey());
            i++;
        }

        StringBuilder builder = new StringBuilder();
        for (int j = 0; j < chars.length; j++) {
            builder.append(encodeMap.get(chars[j]));
        }
        EncodeWrapper wrapper = new EncodeWrapper();
        wrapper.dic = decodeMap;
        wrapper.encodedString = builder.toString();
        return wrapper;
    }

    private String generateString(int i, int size) {
        StringBuilder builder = new StringBuilder();
        if (i == 0) {
            for (int j = i; j < size; j++) {
                builder.append("0");
            }
        } else {
            for (int j = i; j < size - 1; j++) {
                builder.append("0");
            }
            builder.append("1");
        }
        return builder.toString();
    }

    public String decode(String source, HashMap<String, Character> dic) {

        StringBuilder result = new StringBuilder();

        char[] chars = source.toCharArray();
        for (int i = 0; i < chars.length; ) {
            int j = i;
            StringBuilder builder = new StringBuilder();
            // loop invariant: [i,j]
            while (j < chars.length && !dic.containsKey(builder.toString())) {
                builder.append(chars[j]);
                j++;
            }
            if (j <= chars.length && dic.containsKey(builder.toString())) {
                result.append(dic.get(builder.toString()));
                i = j;
            } else {
                return "-1";
            }
        }
        return result.toString();
    }
}

class EncodeWrapper {

    String encodedString;

    HashMap<String, Character> dic;
}

class MapUtil {
    public static <K, V extends Comparable<? super V>> Map<K, V> sortByValue(Map<K, V> map) {
        List<Map.Entry<K, V>> list = new ArrayList<>(map.entrySet());
        // from small to large
        list.sort(Map.Entry.comparingByValue());

        Map<K, V> result = new LinkedHashMap<>();
        for (Map.Entry<K, V> entry : list) {
            result.put(entry.getKey(), entry.getValue());
        }

        return result;
    }
}
